import pandas as pd
from mlxtend.preprocessing import TransactionEncoder
from mlxtend.frequent_patterns import apriori, association_rules

df = pd.read_excel('./餐厅数据.xlsx')

trasations = df['菜品'].str.split(',').to_list()

# 标准化
te = TransactionEncoder()
te_ary = te.fit(trasations).transform(trasations)
de_enecoded = pd.DataFrame(te_ary, columns=te.columns_)

# 使用apriori进行分析
frequent_itemsets = apriori(de_enecoded, min_support=0.1, use_colnames=True)
frequent_itemsets.sort_values(by='support', ascending=False, inplace=True)

# 选额二项集查看
print(frequent_itemsets[frequent_itemsets.itemsets.apply(lambda x:len(x) == 2)])

# 生成关联规则
rules = association_rules(frequent_itemsets, metric='confidence', min_threshold=0.1)

rules = rules.sort_values(by=['lift'], ascending=False)

# 保存模型
frequent_itemsets.to_pickle('frequent_itemsets.pkl')
rules.to_pickle('rules.pkl')